#!/usr/bin/env python
# encoding: utf-8

import argparse
import os
from sklearn.utils import shuffle

parser = argparse.ArgumentParser()
parser.add_argument('--thchs30_test_dir', help='data list', type=str, default="data/tmp_list")
parser.add_argument('--output_trial_path', help='embedding dim', type=str, default="t.lst")
args = parser.parse_args()

path = []
label = []
for spk in os.listdir(args.thchs30_test_dir):
    spk_path = os.path.join(args.thchs30_test_dir, spk)
    for wav in os.listdir(spk_path):
        wav = os.path.join(spk_path, wav)
        label.append(spk)
        path.append(wav)

path, label = shuffle(path, label)

enroll_path = path[:1000]
enroll_label = label[:1000]
test_path = path[1000:]
test_label = label[1000:]

f =  open(args.output_trial_path, "a+")
for x in range(len(enroll_path)):
    for y in range(len(test_path)):
        if enroll_label[x] == test_label[y]:
            f.write("1 {} {}\n".format(enroll_path[x], test_path[y]))
        else:
            f.write("0 {} {}\n".format(enroll_path[x], test_path[y]))

